from solvers.solverbase import SolverBase
from solvers.analytical.analyticalbase import AnalyticalBase
from dolfin import *

class LeapFrog(SolverBase,AnalyticalBase):
    'Leap-frog scheme.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        AnalyticalBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)  #inital condition for backward-euler step to solve for phi^n
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)

        (ib,ib_value,bParts) = problem.get_inflowBCS()        
        if problem.weakBC:
            a = phi*varphi*dx - dt*(inner(nabla_grad(varphi),v*phi)*dx + nabla_div(v)*phi*varphi*dx)
        else:
            n = FacetNormal(mesh)
            a = phi*varphi*dx + dt*(phi*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),v*phi)*dx - nabla_div(v)*phi*varphi*dx) 

        l = phi_*varphi*dx

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
       
        A = assemble(a,exterior_facet_domains=bParts) #A unlike L has ds(0)
        phi = Function(V)
   
        t = float(dt)
        count = 1

        #Backward Euluer step
        phi__ = Function(V)
        L = assemble(l) 
        ibc = DirichletBC(V,ib_value(t),ib)
        ibc.apply(A,L)            
        solve(A,phi.vector(),L)
        
        phi__.assign(phi_)   # hold the initial condition phi_0
        phi_.assign(phi)     # hold the current solution  phi_1
        self.update(problem,phi_,t,float(dt),count)
        t += float(dt)
        count += 1

        #now OK to solve for phi_2 (phi_(n_1))
        #make phi a trialfunction again and define forms
        phi = TrialFunction(V)
        a = phi*varphi*dx

        if problem.weakBC:
            l = phi__*varphi*dx + 2.*dt*(inner(nabla_grad(varphi),phi_*v)*dx + phi_*varphi*nabla_div(v)*dx) 
        else:
            n = FacetNormal(mesh)
            l = phi__*varphi*dx - 2.*dt*(phi_*varphi*inner(v,n)*ds(0) - inner(nabla_grad(varphi),phi_*v)*dx - phi_*varphi*nabla_div(v)*dx) 

        #make phi a function and solve, don't forget assignments
        phi = Function(V)
        while t <= T:
            L = assemble(l,exterior_facet_domains=bParts)  # v is not f(t), if it were, A has to be reassembled as well as L
            ibc = DirichletBC(V,ib_value(t),ib)
            ibc.apply(A,L)
            solve(A,phi.vector(),L)
            phi__.assign(phi_)
            phi_.assign(phi)
            
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1

        problem.save_data(self.saveDir,str(self.CFL))

